import os
import json
import asyncio
import time
from fastapi import FastAPI
from typing import Annotated, Literal
from typing_extensions import TypedDict
from pydantic import BaseModel, Field

from conf.config import settings
from common.logger import logger
# from common import postgres_manager, get_pg_pool
# from common import redis_manager, get_redis_client, get_redis_pool

from langchain_openai import ChatOpenAI
from langchain_core.messages import HumanMessage, AIMessage, ToolMessage


from langgraph.graph import StateGraph, START, END
from langgraph.graph.message import add_messages
from langgraph.prebuilt import ToolNode, tools_condition
from langgraph.checkpoint.memory import MemorySaver
from langgraph.checkpoint.postgres.aio import AsyncPostgresSaver
from langgraph.types import interrupt, Command
from langchain_qwq import ChatQwQ

from agents.mcps.oa.get_psndoc import get_personal_data
from agents.mcps.oa.get_cal_has_sumhour import get_duration_vacation
from agents.mcps.oa.get_kq import get_dong_yang_kq
from agents.mcps.oa.post_qianka import qianka_get_kq, qianka_post
from agents.memory.pgdb_pool import init_postgres_pool
from agents.scipt.join_states import join_state
from agents.prompt.prompts import HOLIDAYS_DURATION_PROMPT, KAOQIN_CHATBOT_PROMPT, PERSONAL_DATA_PROMPT, QIANKA_PROMPT, COMMON_CHATBOT_PROMPT, QIANKA_PROMPT

# class ChatModel():
#     def __init__(self, model_name: str = "qwen32b", is_online_search: bool = False, is_thinking: bool = False):
#         if model_name == "qwen32b":
#             self.model = ChatQwQ(
#                 model=settings.QWEN32_MODEL,
#                 api_key=settings.QWEN32_API_KEY,
#                 base_url=settings.QWEN32_BASE_URL,
#                 streaming=True,
#             )

#         elif model_name == "qwen235b":
#             self.model = ChatQwQ(
#                 model=settings.QWEN235_MODEL,
#                 api_key=settings.QWEN235_API_KEY,
#                 base_url=settings.QWEN235_BASE_URL,
#                 streaming=True,
#             )

#         elif model_name == "deepseek-r1":
#             self.model = ChatOpenAI(
#                 model=settings.QWEN235_MODEL,
#                 api_key=settings.QWEN235_API_KEY,
#                 base_url=settings.QWEN235_BASE_URL,
#                 streaming=True,
#             )
#         else:
#             logger.info("没有找到指定模型, 默认使用Qwen235b")
#             self.model = ChatQwQ(
#                 model=settings.QWEN235_MODEL,
#                 api_key=settings.QWEN235_API_KEY,
#                 base_url=settings.QWEN235_BASE_URL,
#                 streaming=True,
#             )
            

class State(TypedDict):
    """状态类定义"""
    messages: Annotated[list, add_messages]

class LangGraphMCPAgent:
    """LangGraph MCP Agent主类"""

    def __init__(self, model_name: str = "qwen32b", app: FastAPI = None, online_search: bool = False, is_thinking: bool = True):
        """初始化Agent"""
        self.app = app
        self.model = None
        self.model_name = model_name
        self.online_search = online_search
        self.is_thinking = is_thinking
        self.local_qwen3 = None
        self.local_qwen3_with_tools = None
        self.tools = None
        self.graph = None
        self.graph_builder = None
        self.qianka_key = None
        self._setup_chat_model()
        self._setup_tools()

    def _setup_chat_model(self):
        """设置聊天模型"""
        # 定义模型映射表
        model_map = {
            "qwen32b": {
                "model": settings.QWEN32_MODEL,
                "api_key": settings.QWEN32_API_KEY,
                "base_url": settings.QWEN32_BASE_URL,
            },
            "qwen235b": {
                "model": settings.QWEN235_MODEL,
                "api_key": settings.QWEN235_API_KEY,
                "base_url": settings.QWEN235_BASE_URL,
            },
            "deepseek-r1": {
                "model": settings.QWEN235_MODEL,
                "api_key": settings.QWEN235_API_KEY,
                "base_url": settings.QWEN235_BASE_URL,
            },
            "gpt4": {
                "model": settings.QWEN235_MODEL,
                "api_key": settings.QWEN235_API_KEY,
                "base_url": settings.QWEN235_BASE_URL,
            },
        }

        # 获取模型配置，默认使用 qwen235b
        config = model_map.get(self.model_name, model_map["qwen235b"])
        if self.model_name not in model_map:
            logger.info("没有找到指定模型, 默认使用 Qwen235b")

        # 根据 is_thinking 选择模型类
        if self.is_thinking:
            self.model = ChatQwQ(
                model=config["model"],
                api_key=config["api_key"],
                base_url=config["base_url"],
                streaming=True,
                temperature=0.6, 
                top_p=0.95,  
            )
        else:
            self.model = ChatOpenAI(
                model=config["model"],
                api_key=config["api_key"],
                base_url=config["base_url"],
                streaming=True,
                temperature=0.7,
                top_p=0.8,
            )

        # 保持 local_qwen3 的兼容
        self.local_qwen3 = self.model

    def _setup_tools(self):
        """设置工具"""
        self.tools = [
            get_personal_data,
            get_duration_vacation,
            get_dong_yang_kq,
            qianka_get_kq
        ]
        self.local_qwen3_with_tools = self.local_qwen3.bind_tools(self.tools)

    async def chatbot(self, state: State):
        """基础聊天机器人节点"""
        invoke_messages= [{"role": "system","content": COMMON_CHATBOT_PROMPT}]
        for i in state["messages"]:
            if isinstance(i, HumanMessage):
                invoke_messages.append({"role": "user", "content": i.content})
            elif isinstance(i, ToolMessage):
                invoke_messages.append({"role": "tool", "content": i.content, "tool_call_id": i.tool_call_id})
            elif isinstance(i, AIMessage):
                invoke_messages.append({"role": "assistant", "content": i.content})
        response = await self.local_qwen3_with_tools.ainvoke(invoke_messages)
        return {"messages": [response]}

    async def search_holidays_duration_chatbot(self, state: State):
        """处理年假查询的聊天机器人节点"""
        last_message = state["messages"][-1]
        all_chunks = []
        logger.info(f"调用年假函数 search_holidays_duration_chatbot {last_message}")

        async for chunk in self.local_qwen3_with_tools.astream([
            {
                "role": "system",
                "content": HOLIDAYS_DURATION_PROMPT
            },
            {"role": "user", "content": last_message.content}
        ]):
            all_chunks.append(chunk.content)
            aim_message = join_state(chunk, all_chunks)
            if aim_message != '':
                state['messages'].append(aim_message)
            yield {"messages": [chunk]}

    async def person_data_chatbot(self, state: State):
        """处理个人信息的聊天机器人节点"""
        last_message = state["messages"][-1]
        all_chunks = []
        logger.info(f"调用个人信息函数 person_data_chatbot {last_message}")

        async for chunk in self.local_qwen3_with_tools.astream([
            {
                "role": "system",
                "content": PERSONAL_DATA_PROMPT
            },
            {"role": "user", "content": last_message.content}
        ]):
            all_chunks.append(chunk.content)
            aim_message = join_state(chunk, all_chunks)
            if aim_message != '':
                state['messages'].append(aim_message)
            yield {"messages": [chunk]}

    async def kaoqin_chatbot(self, state: State):
        """处理考勤信息的聊天机器人节点"""
        last_message = state["messages"][-1]
        all_chunks = []
        logger.info(f"调用查询考勤记录函数 kaoqin_chatbot {last_message}")

        async for chunk in self.local_qwen3_with_tools.astream([
            {
                "role": "system",
                "content": KAOQIN_CHATBOT_PROMPT
            },
            {"role": "tool", "content": last_message.content, "tool_call_id": last_message.tool_call_id}
        ]):
            all_chunks.append(chunk.content)
            aim_message = join_state(chunk, all_chunks)
            if aim_message != '':
                state['messages'].append(aim_message)
            yield {"messages": [chunk]}

    async def qianka_chatbot(self, state: State):
        """处理签卡的聊天机器人节点"""
        last_message = state["messages"][-1]
        user_id = json.loads(last_message.content)[0].get('empNO')
        emp_dates = [item.get('empDate')
                     for item in json.loads(last_message.content)]
        attends = [item.get('attends')
                   for item in json.loads(last_message.content)]

        qianka_args = {
            "user_id": user_id,
            "emp_dates": emp_dates,
            "attends": attends
        }
        self.qianka_key = f'qianka_{user_id}'
        qianka_args_json = json.dumps(qianka_args)
        redis = self.app.state.redis
        await redis.set(self.qianka_key, qianka_args_json)

        all_chunks = []
        logger.info(f"签卡MCP考勤结果 qianka_chatbot {last_message.content}")

        async for chunk in self.local_qwen3_with_tools.astream([
            {
                "role": "system",
                "content": QIANKA_PROMPT
            },
            {"role": "user", "content": last_message.content}
        ]):
            all_chunks.append(chunk.content)
            aim_message = join_state(chunk, all_chunks)
            if aim_message != '':
                state['messages'].append(aim_message)
            yield {"messages": [chunk]}

    async def qianka_interrupt(self, state: State) -> Command[Literal["approved_path", "rejected_path"]]:
        """签卡确认中断节点"""
        is_approved = await asyncio.to_thread(interrupt, "确认签卡吗？(yes or no):")
        if is_approved in ["approve", "yes", "是", "y", "Y", "确认", "同意"]:
            return Command(goto="approved_path")
        else:
            return Command(goto="rejected_path")

    async def classify_tool(self, state: State) -> Literal["chatbot", "search_holidays_duration_chatbot", "person_data_chatbot", "qianka_chatbot", "kaoqin_chatbot"]:
        """工具分类器"""
        last_tool_name = state["messages"][-1].name

        if last_tool_name == "get_duration_vacation":
            return "search_holidays_duration_chatbot"
        elif last_tool_name == "get_personal_data":
            return "person_data_chatbot"
        elif last_tool_name == "qianka_get_kq":
            return "qianka_chatbot"
        elif last_tool_name == "get_dong_yang_kq":
            return "kaoqin_chatbot"
        else:
            return "chatbot"

    async def approved_node(self, state: State) -> State:
        """签卡批准节点"""
        redis = self.app.state.redis
        # 注意：这里需要修复原代码中qianka_key变量作用域的问题
        # 应该从state或redis中获取正确的key
        qianka_data = await redis.get(self.qianka_key)
        if qianka_data:
            qianka_info = json.loads(qianka_data)
            user_id = qianka_info.get('user_id')
            emp_dates = qianka_info.get('emp_dates')
            attends = qianka_info.get('attends')
            # qianka_res = qianka_post(user_id='282696', empdates=['2025-08-04', '2025-08-05'], attends=attends[:2])
            # print(qianka_res)
            # qianka_post(user   # qianka_post(user_id=user_id, empdates=emp_dates, attends=attends)

        return {"messages": [f'✅ 成功提交{emp_dates}的签卡申请']}

    async def rejected_node(self, state: State) -> State:
        """�"签卡拒绝节点"""
        return {"messages": ['❌ 用户已拒绝签卡申请.']}

    async def end_note(self, state: State):
        """结束节点"""
        logger.info(f"结束节点 end_note 输出：{state}")
        # print("===========end_node:===============", "\n", state, "\n")

    def _build_graph(self):
        """构建图结构"""
        self.graph_builder = StateGraph(State)

        # 创建工具节点
        tool_node = ToolNode(tools=self.tools)

        # 添加节点
        self.graph_builder.add_node("chatbot", self.chatbot)
        self.graph_builder.add_node("tools", tool_node)
        self.graph_builder.add_node("search_holidays_duration_chatbot", self.search_holidays_duration_chatbot)
        self.graph_builder.add_node("person_data_chatbot", self.person_data_chatbot)
        self.graph_builder.add_node("kaoqin_chatbot", self.kaoqin_chatbot)
        self.graph_builder.add_node("qianka_chatbot", self.qianka_chatbot)
        self.graph_builder.add_node("qianka_interrupt", self.qianka_interrupt)
        self.graph_builder.add_node("approved_path", self.approved_node)
        self.graph_builder.add_node("rejected_path", self.rejected_node)
        self.graph_builder.add_node("end_note", self.end_note)

        # 添加边
        self.graph_builder.add_edge(START, "chatbot")
        self.graph_builder.add_conditional_edges("chatbot", tools_condition)
        self.graph_builder.add_conditional_edges("tools", self.classify_tool)
        self.graph_builder.add_edge("qianka_chatbot", "qianka_interrupt")
        self.graph_builder.add_edge('search_holidays_duration_chatbot', "end_note")
        self.graph_builder.add_edge('person_data_chatbot', "end_note")
        self.graph_builder.add_edge("kaoqin_chatbot", "end_note")
        self.graph_builder.add_edge("chatbot", "end_note")
        self.graph_builder.add_edge("approved_path", "end_note")
        self.graph_builder.add_edge("rejected_path", "end_note")
        self.graph_builder.add_edge("end_note", END)

    async def initialize(self):
        """初始化Agent，包括构建图和设置检查点"""
        logger.info("调用 LangGraphMCPAgent.initialize 方法")

        self._build_graph()

        try:
            self.graph = self.graph_builder.compile(
                checkpointer=self.app.state.checkpointer,
                store=self.app.state.store)

            logger.info("graph构建成功")
            return self.graph

        except Exception as e:
            logger.error(f"graph构建失败: {e}")
            raise e

    def get_graph(self):
        """获取构建好的图"""
        if self.graph is None:
            raise RuntimeError(
                "Graph not initialized. Please call initialize() first.")
        return self.graph